#include "BottomLevelAccelerationStructure.h"
#include "../../Core/CoreGlobals.h"
#include "../../Utils/LogUtil.h"
#include "../../Templates/NewObject.h"
#include "../RenderState.h"
#include "../../Game/Process/RenderProcess.h"
#include "../../Math/Math.h"

BottomLevelAccelerationStructure::Builder::Builder() noexcept = default;
BottomLevelAccelerationStructure::Builder::~Builder() noexcept = default;
BottomLevelAccelerationStructure::Builder::Builder(BottomLevelAccelerationStructure::Builder const& rhs) noexcept = default;
BottomLevelAccelerationStructure::Builder::Builder(BottomLevelAccelerationStructure::Builder && rhs) noexcept = default;
BottomLevelAccelerationStructure::Builder& BottomLevelAccelerationStructure::Builder::operator=(BottomLevelAccelerationStructure::Builder const& rhs) noexcept = default;
BottomLevelAccelerationStructure::Builder& BottomLevelAccelerationStructure::Builder::operator=(BottomLevelAccelerationStructure::Builder && rhs) noexcept = default;

BottomLevelAccelerationStructure::BottomLevelAccelerationStructure(BuilderDetails & builderDetails)
{
	mType = VK_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL_KHR;
	mFlags = builderDetails.flags;
	mAccelerationStructureInstance = builderDetails.accelerationStructureInstance;
	mGeometries.assign(builderDetails.geometries.begin(), builderDetails.geometries.end());
	mAccelerationStructureBuildSizesInfo = {};
	mAccelerationStructureBuildSizesInfo.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_SIZES_INFO_KHR;
	mIsBuilded = false;

	// Make sure all the geometry same type
	const VkGeometryTypeKHR geometryType = mGeometries[0]->GetAccelerationStructureGeometry().geometryType;
	const uint32 geometryCount = static_cast<uint32>(mGeometries.size());
	mVkGeometries.resize(geometryCount, {});
	mMaxPrimitiveCounts.resize(geometryCount);
	mAccelerationStructureBuildRangeInfos.resize(geometryCount, nullptr);
	for (uint32 i = 0; i < geometryCount; ++i)
	{
		assert(mGeometries[i]->GetAccelerationStructureGeometry().geometryType == geometryType);
		mVkGeometries[i] = mGeometries[i]->GetAccelerationStructureGeometry();
		mMaxPrimitiveCounts[i] = mGeometries[i]->GetPrimitiveCount();
		mAccelerationStructureBuildRangeInfos[i] = mGeometries[i]->GetRangeInfo();
	}

	if (mFlags & VK_BUILD_ACCELERATION_STRUCTURE_ALLOW_COMPACTION_BIT_KHR == VK_BUILD_ACCELERATION_STRUCTURE_ALLOW_COMPACTION_BIT_KHR)
	{
		// wait for the TopLevel to compaction this blas
		return;
	}

	// Get Size Info
	VkAccelerationStructureBuildGeometryInfoKHR accelerationStructureSizeGeometryInfo{};
	accelerationStructureSizeGeometryInfo.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_GEOMETRY_INFO_KHR;
	accelerationStructureSizeGeometryInfo.type = mType;
	accelerationStructureSizeGeometryInfo.flags = mFlags;
	accelerationStructureSizeGeometryInfo.geometryCount = static_cast<uint32>(mVkGeometries.size());
	accelerationStructureSizeGeometryInfo.pGeometries = mVkGeometries.data();

	vkGetAccelerationStructureBuildSizesKHR(
		GVulkanRHI->GetLogicalDevice(),
		VK_ACCELERATION_STRUCTURE_BUILD_TYPE_DEVICE_KHR,
		&accelerationStructureSizeGeometryInfo,
		mMaxPrimitiveCounts.data(),
		&mAccelerationStructureBuildSizesInfo);

	CreateAccelerationStructureBuffer(mAccelerationStructureBuildSizesInfo.accelerationStructureSize, mAccelerationStructureBuffer, mAccelerationStructureMemory);

	VkAccelerationStructureCreateInfoKHR accelerationStructureCreateInfo{};
	accelerationStructureCreateInfo.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_CREATE_INFO_KHR;
	accelerationStructureCreateInfo.buffer = mAccelerationStructureBuffer;
	accelerationStructureCreateInfo.size = mAccelerationStructureBuildSizesInfo.accelerationStructureSize;
	accelerationStructureCreateInfo.type = mType;
	vkCreateAccelerationStructureKHR(GVulkanRHI->GetLogicalDevice(), &accelerationStructureCreateInfo, nullptr, &mAccelerationStructure);

	VkBuffer scratchBuffer;
	VkDeviceMemory scratchMemory;
	VkDeviceAddress scratchAddress;
	CreateScratchBuffer(mAccelerationStructureBuildSizesInfo.buildScratchSize, scratchBuffer, scratchMemory, scratchAddress);

	VkAccelerationStructureBuildGeometryInfoKHR accelerationStructureBuildGeometryInfo{};
	accelerationStructureBuildGeometryInfo.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_GEOMETRY_INFO_KHR;
	accelerationStructureBuildGeometryInfo.type = mType;
	accelerationStructureBuildGeometryInfo.flags = mFlags;
	accelerationStructureBuildGeometryInfo.mode = VK_BUILD_ACCELERATION_STRUCTURE_MODE_BUILD_KHR;
	accelerationStructureBuildGeometryInfo.dstAccelerationStructure = mAccelerationStructure;
	accelerationStructureBuildGeometryInfo.geometryCount = static_cast<uint32>(mVkGeometries.size());
	accelerationStructureBuildGeometryInfo.pGeometries = mVkGeometries.data();
	accelerationStructureBuildGeometryInfo.scratchData.deviceAddress = scratchAddress;

	if (GVulkanRHI->GetAccelerationStructureFeatures().accelerationStructureHostCommands)
	{
		// Implementation supports building acceleration structure building on host
		vkBuildAccelerationStructuresKHR(
			GVulkanRHI->GetLogicalDevice(),
			VK_NULL_HANDLE,
			1,
			&accelerationStructureBuildGeometryInfo,
			mAccelerationStructureBuildRangeInfos.data());
	}
	else
	{
		VkCommandBuffer commandBuffer = GRenderState->BeginSingleCommand();
		vkCmdBuildAccelerationStructuresKHR(
			commandBuffer,
			1,
			&accelerationStructureBuildGeometryInfo,
			mAccelerationStructureBuildRangeInfos.data());
		GRenderState->FlushSingleCommand(commandBuffer);
	}

	vkDestroyBuffer(GVulkanRHI->GetLogicalDevice(), scratchBuffer, nullptr);
	vkFreeMemory(GVulkanRHI->GetLogicalDevice(), scratchMemory, nullptr);

	// Get deviceAddress of as
	GetInstanceAddress();

	mIsBuilded = true;
}

void BottomLevelAccelerationStructure::AddGeometry(RefCountPtr<AccelerationStructureGeometry> newGeometry)
{
	for (auto& geometry : mGeometries)
	{
		assert(geometry != newGeometry);
	}
	mGeometries.push_back(newGeometry);
}

void BottomLevelAccelerationStructure::RemoveGeometry(RefCountPtr<AccelerationStructureGeometry> oldGeometry)
{
	for (uint32 i = 0; i < static_cast<uint32>(mGeometries.size()); ++i)
	{
		if (mGeometries[i] == oldGeometry)
		{
			mGeometries.erase(mGeometries.begin() + i);
			return;
		}
	}
	assert(0);
}

void BottomLevelAccelerationStructure::UpdateTransform(VkTransformMatrixKHR transform)
{
	mAccelerationStructureInstance.transform = transform;
}

void BottomLevelAccelerationStructure::UpdateBlas(VkBuildAccelerationStructureFlagsKHR flags)
{
	// make sure is build already
	assert(mIsBuilded);

	// Make sure all the geometry same type
	const VkGeometryTypeKHR geometryType = mGeometries[0]->GetAccelerationStructureGeometry().geometryType;
	const uint32 geometryCount = static_cast<uint32>(mGeometries.size());
	mVkGeometries.resize(geometryCount, {});
	mMaxPrimitiveCounts.resize(geometryCount);
	mAccelerationStructureBuildRangeInfos.resize(geometryCount, nullptr);
	for (uint32 i = 0; i < geometryCount; ++i)
	{
		assert(mGeometries[i]->GetAccelerationStructureGeometry().geometryType == geometryType);
		mVkGeometries[i] = mGeometries[i]->GetAccelerationStructureGeometry();
		mMaxPrimitiveCounts[i] = mGeometries[i]->GetPrimitiveCount();
		mAccelerationStructureBuildRangeInfos[i] = mGeometries[i]->GetRangeInfo();
	}

	VkAccelerationStructureBuildGeometryInfoKHR accelerationStructureBuildGeometryInfo{};
	accelerationStructureBuildGeometryInfo.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_GEOMETRY_INFO_KHR;
	accelerationStructureBuildGeometryInfo.type = mType;
	accelerationStructureBuildGeometryInfo.flags = flags == VK_BUILD_ACCELERATION_STRUCTURE_FLAG_BITS_MAX_ENUM_KHR ? mFlags : VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_TRACE_BIT_KHR;
	accelerationStructureBuildGeometryInfo.mode = VK_BUILD_ACCELERATION_STRUCTURE_MODE_UPDATE_KHR;
	accelerationStructureBuildGeometryInfo.geometryCount = static_cast<uint32>(mVkGeometries.size());
	accelerationStructureBuildGeometryInfo.pGeometries = mVkGeometries.data();
	accelerationStructureBuildGeometryInfo.srcAccelerationStructure = mAccelerationStructure;
	accelerationStructureBuildGeometryInfo.dstAccelerationStructure = mAccelerationStructure;

	// Find size to build on the device
	vkGetAccelerationStructureBuildSizesKHR(
		GVulkanRHI->GetLogicalDevice(),
		VK_ACCELERATION_STRUCTURE_BUILD_TYPE_DEVICE_KHR,
		&accelerationStructureBuildGeometryInfo,
		mMaxPrimitiveCounts.data(),
		&mAccelerationStructureBuildSizesInfo);

	VkBuffer scratchBuffer;
	VkDeviceMemory scratchMemory;
	VkDeviceAddress scratchAddress;
	CreateScratchBuffer(mAccelerationStructureBuildSizesInfo.buildScratchSize, scratchBuffer, scratchMemory, scratchAddress);

	accelerationStructureBuildGeometryInfo.scratchData.deviceAddress = scratchAddress;

	if (GVulkanRHI->GetAccelerationStructureFeatures().accelerationStructureHostCommands)
	{
		// Implementation supports building acceleration structure building on host
		vkBuildAccelerationStructuresKHR(
			GVulkanRHI->GetLogicalDevice(),
			VK_NULL_HANDLE,
			1,
			&accelerationStructureBuildGeometryInfo,
			mAccelerationStructureBuildRangeInfos.data());
	}
	else
	{
		VkCommandBuffer commandBuffer = GRenderState->BeginSingleCommand();
		vkCmdBuildAccelerationStructuresKHR(
			commandBuffer,
			1,
			&accelerationStructureBuildGeometryInfo,
			mAccelerationStructureBuildRangeInfos.data());
		GRenderState->FlushSingleCommand(commandBuffer);
	}

	vkDestroyBuffer(GVulkanRHI->GetLogicalDevice(), scratchBuffer, nullptr);
	vkFreeMemory(GVulkanRHI->GetLogicalDevice(), scratchMemory, nullptr);

	// Get deviceAddress of as
	GetInstanceAddress();
}

VkAccelerationStructureInstanceKHR& BottomLevelAccelerationStructure::GetAccelerationStructureInstance()
{
	return mAccelerationStructureInstance;
}

bool BottomLevelAccelerationStructure::IsBuilded() const
{
	return mIsBuilded;
}

std::vector<VkAccelerationStructureGeometryKHR>& BottomLevelAccelerationStructure::GetVkGeometries()
{
	return mVkGeometries;
}

std::vector<uint32>& BottomLevelAccelerationStructure::GetMaxPrimitiveCounts()
{
	return mMaxPrimitiveCounts;
}

std::vector<VkAccelerationStructureBuildRangeInfoKHR*>& BottomLevelAccelerationStructure::GetRangeInfos()
{
	return mAccelerationStructureBuildRangeInfos;
}

VkAccelerationStructureBuildSizesInfoKHR& BottomLevelAccelerationStructure::GetSizeInfo()
{
	return mAccelerationStructureBuildSizesInfo;
}

VkAccelerationStructureKHR& BottomLevelAccelerationStructure::GetClearStructure()
{
	return mClearStructure;
}

void BottomLevelAccelerationStructure::SaveClearStructure()
{
	mClearStructureBuffer = mAccelerationStructureBuffer;
	mClearStructureMemory = mAccelerationStructureMemory;
	mClearStructure = mAccelerationStructure;
}

void BottomLevelAccelerationStructure::DeleteClearStructure()
{
	vkDestroyAccelerationStructureKHR(GVulkanRHI->GetLogicalDevice(), mClearStructure, nullptr);
	vkDestroyBuffer(GVulkanRHI->GetLogicalDevice(), mClearStructureBuffer, nullptr);
	vkFreeMemory(GVulkanRHI->GetLogicalDevice(), mClearStructureMemory, nullptr);
}

void BottomLevelAccelerationStructure::GetInstanceAddress()
{
	// Get deviceAddress of as
	VkAccelerationStructureDeviceAddressInfoKHR accelerationStructureDeviceAddressInfo{};
	accelerationStructureDeviceAddressInfo.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_DEVICE_ADDRESS_INFO_KHR;
	accelerationStructureDeviceAddressInfo.accelerationStructure = mAccelerationStructure;
	mAccelerationStructureInstance.accelerationStructureReference = vkGetAccelerationStructureDeviceAddressKHR(GVulkanRHI->GetLogicalDevice(), &accelerationStructureDeviceAddressInfo);
	assert(mAccelerationStructureInstance.accelerationStructureReference != 0);

	mIsBuilded = true;
}

BottomLevelAccelerationStructure::Builder& BottomLevelAccelerationStructure::Builder::SetFlags(VkBuildAccelerationStructureFlagsKHR flags) noexcept
{
	mImpl->flags = flags;
	return *this;
}

BottomLevelAccelerationStructure::Builder& BottomLevelAccelerationStructure::Builder::SetTransform(glm::mat4 transform) noexcept
{
	// glm::mat4 modelMatrix = glm::transpose(transform);
	// memcpy(&mImpl->accelerationStructureInstance.transform, &modelMatrix, sizeof(VkTransformMatrixKHR));

	mImpl->accelerationStructureInstance.transform.matrix[0][0] = transform[0][0];
	mImpl->accelerationStructureInstance.transform.matrix[0][1] = transform[1][0];
	mImpl->accelerationStructureInstance.transform.matrix[0][2] = transform[2][0];
	mImpl->accelerationStructureInstance.transform.matrix[0][3] = transform[3][0];
	mImpl->accelerationStructureInstance.transform.matrix[1][0] = transform[0][1];
	mImpl->accelerationStructureInstance.transform.matrix[1][1] = transform[1][1];
	mImpl->accelerationStructureInstance.transform.matrix[1][2] = transform[2][1];
	mImpl->accelerationStructureInstance.transform.matrix[1][3] = transform[3][1];
	mImpl->accelerationStructureInstance.transform.matrix[2][0] = transform[0][2];
	mImpl->accelerationStructureInstance.transform.matrix[2][1] = transform[1][2];
	mImpl->accelerationStructureInstance.transform.matrix[2][2] = transform[2][2];
	mImpl->accelerationStructureInstance.transform.matrix[2][3] = transform[3][2];


	return *this;
}

BottomLevelAccelerationStructure::Builder& BottomLevelAccelerationStructure::Builder::SetInstanceCustomIndex(uint32 instanceCustomIndex) noexcept
{
	mImpl->accelerationStructureInstance.instanceCustomIndex = instanceCustomIndex;
	return *this;
}

BottomLevelAccelerationStructure::Builder& BottomLevelAccelerationStructure::Builder::SetMask(uint32 mask) noexcept
{
	mImpl->accelerationStructureInstance.mask = mask;
	return *this;
}

BottomLevelAccelerationStructure::Builder& BottomLevelAccelerationStructure::Builder::SetInstanceShaderBindingTableRecordOffset(uint32 instanceShaderBindingTableRecordOffset) noexcept
{
	mImpl->accelerationStructureInstance.instanceShaderBindingTableRecordOffset = instanceShaderBindingTableRecordOffset;
	return *this;
}

BottomLevelAccelerationStructure::Builder& BottomLevelAccelerationStructure::Builder::SetInstanceFlags(VkGeometryInstanceFlagsKHR flags) noexcept
{
	mImpl->accelerationStructureInstance.flags = flags;
	return *this;
}

BottomLevelAccelerationStructure::Builder& BottomLevelAccelerationStructure::Builder::AddGeometry(RefCountPtr<AccelerationStructureGeometry> geometry) noexcept
{
	mImpl->geometries.push_back(geometry);
	return *this;
}

RefCountPtr<BottomLevelAccelerationStructure> BottomLevelAccelerationStructure::Builder::Build()
{
	assert(IsInRenderThread());
	return RefCountPtr<BottomLevelAccelerationStructure>(NewObject<BottomLevelAccelerationStructure>(*mImpl));
}

RefCountPtr<BottomLevelAccelerationStructure> BottomLevelAccelerationStructure::Builder::Build(RenderProcess * process, const std::string & name)
{
	assert(IsInRenderThread());
	BottomLevelAccelerationStructure* bottomLevelAccelerationStructure = NewObject<BottomLevelAccelerationStructure>(*mImpl);
	process->AddAccelerationStructure(name, bottomLevelAccelerationStructure);
	return RefCountPtr<BottomLevelAccelerationStructure>(bottomLevelAccelerationStructure);
}

void BottomLevelAccelerationStructure::ReleaseResource()
{
	if (IsInRenderThread())
	{
		if (mProcess != nullptr)
		{
			mProcess->RemoveAccelerationStructure(mName, this);
		}
		DeleteObject<BottomLevelAccelerationStructure>(this);
	}
	else
	{
		AsyncTask(GRenderThreadId, [this]() {
			if (mProcess != nullptr)
			{
				mProcess->RemoveAccelerationStructure(mName, this);
			}
			DeleteObject<BottomLevelAccelerationStructure>(this);
			});
	}
}

